package com.tmian.aiagent.SuperAgent;

import cn.hutool.core.collection.CollUtil;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.tmian.aiagent.SuperAgent.model.AgentState;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;

import java.util.List;
import java.util.stream.Collectors;

/**
 * 处理工具调用的基础代理类，具体实现了 think 和 act 方法，可以用作创建实例的父类
 */
@EqualsAndHashCode(callSuper = true)
@Data
@Slf4j
public class ToolCallAgent extends ReActAgent{

    //可用工具
    private final ToolCallback[] availabletools;

    //工具响应
    private ChatResponse toolCallChatResponse;

    //工具管理
    private final ToolCallingManager toolCallingManager;

    //禁用内置的工具调用机制，自己管理工具调用
    private final ChatOptions chatOptions;

    public ToolCallAgent(ToolCallback[] toolCallbacks){
        this.availabletools = toolCallbacks;
        this.toolCallingManager = ToolCallingManager.builder().build();
        this.chatOptions = DashScopeChatOptions
                .builder()
                .withProxyToolCalls(true)
                .build();

    }
    /**
     * 处理当前状态并决定下一步行动
     *
     * @return 是否需要执行行动
     */
    @Override
    public boolean think() {
        if(getNextStepPrompt() != null && !getNextStepPrompt().isEmpty()){
            getMessageList().add(new UserMessage(getNextStepPrompt()));
        }

        //获取到对话上下文，用于这次对话的用户提示词
        List<Message> messagelist = getMessageList();

        Prompt prompt = new Prompt(messagelist, chatOptions);

        try {
            //调用大模型，获取响应结果
            ChatResponse chatResponse = getChatClient().prompt(prompt)
                    .system(getSystemPrompt())
                    .tools(availabletools)
                    .call()
                    .chatResponse();

            //记录响应，用于act
            this.toolCallChatResponse = chatResponse;
            AssistantMessage assistantMessage = chatResponse.getResult().getOutput();

            String result = assistantMessage.getText();
            //获取大模型要调用的工具列表
            List<AssistantMessage.ToolCall> toolCallList = assistantMessage.getToolCalls();

            log.info(getName() + "的思考 : " + result);
            log.info(getName() + "选择了 " + toolCallList.size() + " 个工具来使用");

            //获取工具名称及参数
            String tollCall = toolCallList.stream()
                    .map(toolCall -> String.format("工具名称: %s, 参数: %s",
                            toolCall.name(),
                            toolCall.arguments()))
                    .collect(Collectors.joining("\n"));

            log.info(tollCall);

            if(toolCallList.isEmpty()){
                //无工具调用，添加助手消息
                getMessageList().add(assistantMessage);
                return false;
            }else{
                //有工具调用，无需添加助手消息，调用工具时会自动添加
                return true;
            }
        } catch (Exception e) {
            log.error(getName() + "的思考过程遇到了问题: " + e.getMessage());
            getMessageList().add(new AssistantMessage("处理时遇到错误" + e.getMessage()));
            return false;
        }
    }

    /**
     * 执行工具调用并处理结果
     *
     * @return 执行结果
     */
    @Override
    public String act() {
        if(!toolCallChatResponse.hasToolCalls()){
            return "无工具调用";
        }
        //调用工具管理器执行工具
        Prompt prompt = new Prompt(getMessageList(), chatOptions);
        ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, toolCallChatResponse);

        //记录上下文，conversationHistory()中包含了工具调用及工具调用的结果
        setMessageList(toolExecutionResult.conversationHistory());

        //获取工具调用的结果
        ToolResponseMessage toolResponseMessage = (ToolResponseMessage) CollUtil.getLast(toolExecutionResult.conversationHistory());

        String results = toolResponseMessage.getResponses().stream()
                .map(response -> "工具" + response.name() + "完成了它的任务" + response.responseData())
                .collect(Collectors.joining("\n"));

        //判断是否调用了终止工具
        boolean match = toolResponseMessage.getResponses()
                .stream().anyMatch(response -> "doTerminate".equals(response.name()));
        if(match){
            setState(AgentState.FINISHED);
        }

        log.info(results);
        return results;
    }

    @Override
    protected void cleanup() {

    }
}
